import matplotlib.pyplot  as plt
from basic import load_raindata
import pandas as pd
import random

def load_data(indir):
    df = pd.read_csv(indir + "/bd2019-weather-prediction-training-20190608.csv" )
    return df


def rain_relation(df):
    cities = list(df.groupby(by=['city', 'station'])['city', 'station'].count().index)
    rows = 4
    cols = 3
    f, ax = plt.subplots(rows, cols)
    ij = 0

    for city,station in cities:
        if random.random() > 0.17:
            continue
        print(city, station)
        i = int(ij / cols)
        j = ij % cols

        v = df[(df['rain08'] > 0) & (df['rain20'] > 0) & (df['rain20'] < 999990) & (df['rain08'] < 999990) & (df['city'] == city) & (df['station'] == station)]
        vv = v.sort_values(by = 'rain20')
        x = vv['rain20']
        y = vv['rain08']

        ax[i][j].scatter(x, y)
        ij += 1

        print(city, len(v))
    plt.show()


def show_fields_days(df, field):
    cities = list(df.groupby(by=['city', 'station'])['city', 'station'].count().index)
    rows = 4
    cols = 3
    f, ax = plt.subplots(rows, cols)
    ij = 0

    for city,station in cities:
        if random.random() > 0.16 or ij > 11:
            continue
        i = int(ij / cols)
        j = ij % cols

        v = df[(df[field] < 999990)  & (df['city'] == city) & (df['station'] == station)]
        print(city, station, len(v), len(df[(df['city'] == city) & (df['station'] == station)]) )
        y = v.sort_values(by = 'date')[field]
        x = range(1, len(y) + 1)

        ax[i][j].scatter(x, y)
        ij += 1
    plt.show()


def field_days_scatter(df, field_name, city, station, ax, i ,j ):
    v = df[(df[field_name] < 999990) & (df['city'] == city) & (df['station'] == station)]
    print(city, station, ":",len(v), len(df[(df['city'] == city) & (df['station'] == station)]))
    y = v.sort_values(by='date')[field_name]
    x = range(1, len(y) + 1)
    ax[i][j].scatter(x, y)

def field_value_hist(df, field_name, bins, city, station, ax, i ,j ):
    v = df[(df[field_name] < 999990) & (df['city'] == city) & (df['station'] == station)]
    print(city, station,  ":", len(v), len(df[(df['city'] == city) & (df['station'] == station)]))
    y = v[field_name]
    ax[i][j].hist( y, bins= bins)

def run_show(df, kwargs = {}):
    cities = list(df.groupby(by=['city', 'station'])['city', 'station'].count().index)
    rows = 4
    cols = 3
    f, ax = plt.subplots(rows, cols)
    ij = 0

    for city, station in cities:
        if random.random() > 0.16 or ij > 11:
            continue

        i = int(ij / cols)
        j = ij % cols
        if kwargs.get('plot') == 'field_days':
            field_days_scatter(df, kwargs.get('field'), city, station, ax, i, j)
        if kwargs.get('plot') == 'field_hist':
            field_value_hist(df, kwargs.get('field'), kwargs.get('bins'), city, station, ax, i, j)

        ij += 1
    plt.show()


def fill_rain_cloud(df):
    grp = df.groupby(['city', 'date'])
    for k, gdata in grp:

        print("run:{}".format(k), end= ">")
        fixes = 0
        if len(gdata[(gdata['rain20'] == 999990) | (gdata['rain08'] == 999990)]) > 0:

            rain20miss = gdata[(gdata['rain20'] == 999990) & (gdata['rain08'] != 999990)]
            rain08miss = gdata[(gdata['rain20'] != 999990) & (gdata['rain08'] == 999990)]
            if len(rain20miss) > 0:
                gdata.loc[(gdata['rain20'] == 999990) & (gdata['rain08'] != 999990), 'rain20'] = rain20miss['rain08']
            if len(rain08miss) > 0:
                gdata.loc[(gdata['rain20'] != 999990) & (gdata['rain08'] == 999990), 'rain08'] = rain08miss['rain20']

            rain20s = gdata[(gdata['rain20'] != 999990)]
            rain08s = gdata[(gdata['rain08'] != 999990)]

            if len(rain20s) > 0:
                mean = rain20s['rain20'].mean()
                gdata.loc[(gdata['rain20'] == 999990), 'rain20'] = (0 if mean < 0.083 else mean )
            if len(rain08s) > 0:
                mean = rain08s['rain08'].mean()
                gdata.loc[(gdata['rain08'] == 999990), 'rain08'] = (0 if mean < 0.083 else mean )
            df.loc[(df['city'] == k[0]) & (df['date'] == k[1]), ['rain20', 'rain08']] = gdata[['rain20', 'rain08']]
            print(" rain fixed  ", end="")
            fixes += 1
        if len(gdata[(gdata['cloud'] >= 999990)]) > 0:
            clouds = gdata[(gdata['cloud'] < 999990)]
            if len(clouds) > 0:
                gdata.loc[(gdata['cloud'] >= 999990), 'cloud'] = clouds['cloud'].mean()
                df.loc[(df['city'] == k[0]) & (df['date'] == k[1]), ['cloud']] = gdata[['cloud']]
                print(" cloud fixed  ",end="")
            else:
                print(" cloud failed  ", end="")
            fixes += 1
        if fixes == 0:
            print(" no need")
        else:
            print()

if __name__ == "__main__":
    df = load_raindata("fix_rain.csv")
   # df = load_data("d:/sfxy191")
  #  fill_rain_cloud(df)
   # df.sort_values(['city', 'date', 'station']).to_csv("./fix_rain.csv",index=0)
   # df1.to_csv("./train.csv",index=0)
  #  rain_relation(df)
  #  show_fields_days(df, 'wind_speed')
   # rain_bins = list(map(lambda x:x/10, range(0,10))) + \
   #             list(map(lambda x: 1.0 + x / 10, range(1, 8,2))) + \
             #   list(map(lambda x: 2.0 + x / 10, range(1, 9, 5))) + [3, 4, 5.5, 7.7, 10, 15, 20, 50, 100]
    run_show(df, {"field":'visibility', 'plot':'field_days', 'bins':20})